import torchvision.transforms as transforms
import matplotlib.pyplot as plt
from PIL import Image

def processImage():
    # Define the transformation
    transform_pipeline = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.GaussianBlur(kernel_size=(15, 15), sigma=(5.0, 5.0)),  # 模糊处理
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
    ])

    # Load a sample image
    image_path = r'./images/02095.jpg'  # Replace with your image path
    original_image = Image.open(image_path)

    # Apply transformations step-by-step
    resize_transform = transforms.Resize((224, 224))
    blur_transform = transforms.GaussianBlur(kernel_size=(31, 31), sigma=(10.0, 10.0))
    to_tensor_transform = transforms.ToTensor()
    normalize_transform = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])

    resized_image = resize_transform(original_image)
    blurred_image = blur_transform(resized_image)
    tensor_image = to_tensor_transform(blurred_image)
    normalized_image = normalize_transform(tensor_image)

    # Plot each stage
    fig, axes = plt.subplots(1, 4, figsize=(16, 8))

    axes[0].imshow(original_image)
    axes[0].set_title("Original Image")
    axes[0].axis("off")

    axes[1].imshow(resized_image)
    axes[1].set_title("Resized Image")
    axes[1].axis("off")

    axes[2].imshow(blurred_image)
    axes[2].set_title("Blurred Image")
    axes[2].axis("off")

    # Convert normalized tensor back to image for visualization
    denormalized_image = (normalized_image * 0.5 + 0.5).permute(1, 2, 0).numpy()
    axes[3].imshow(denormalized_image)
    axes[3].set_title("Normalized Image")
    axes[3].axis("off")

    plt.tight_layout()
    plt.show()

if __name__ == '__main__':
    processImage()